#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
端口映射工具
功能：将本地127.0.0.1绑定的指定端口映射到其他IP能访问到的端口
作者：HonghaiProxy
"""

import socket
import threading
import sys
import argparse
import logging
import tkinter as tk
from tkinter import ttk, scrolledtext
import queue
import time
import json
import os.path

# 尝试导入pystray库来实现系统托盘功能
try:
    from pystray import Icon as TrayIcon, Menu as TrayMenu, MenuItem as TrayMenuItem
    from PIL import Image, ImageDraw
    HAS_PYSTRAY = True
except ImportError:
    print("警告: 无法导入pystray和PIL库，系统托盘功能将不可用。")
    print("请运行: pip install pystray pillow 来安装必要的依赖。")
    HAS_PYSTRAY = False

# 自定义日志处理器，用于将日志显示到GUI界面
class QueueHandler(logging.Handler):
    def __init__(self, log_queue):
        super().__init__()
        self.log_queue = log_queue
    
    def emit(self, record):
        self.log_queue.put(self.format(record))

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class PortForwarder:
    def __init__(self, local_ip, local_port, remote_ip, remote_port, buffer_size=4096):
        """初始化端口转发器
        
        Args:
            local_ip: 本地监听IP
            local_port: 本地监听端口
            remote_ip: 远程目标IP（这里是127.0.0.1）
            remote_port: 远程目标端口
            buffer_size: 缓冲区大小
        """
        self.local_ip = local_ip
        self.local_port = local_port
        self.remote_ip = remote_ip
        self.remote_port = remote_port
        self.buffer_size = buffer_size
        self.server_socket = None
        self.threads = []
        self.running = False
        self.server_thread = None
    
    def start(self):
        """启动转发服务器"""
        try:
            # 创建并配置服务器套接字
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            
            # 绑定到本地IP和端口，0.0.0.0表示所有网络接口都可访问
            self.server_socket.bind((self.local_ip, self.local_port))
            self.server_socket.listen(5)
            
            self.running = True
            logger.info(f"端口转发已启动: {self.local_ip}:{self.local_port} -> {self.remote_ip}:{self.remote_port}")
            
            # 主循环：接受新连接
            try:
                while self.running:
                    # 检查套接字是否仍然有效
                    if self.server_socket is None:
                        break
                    
                    try:
                        # 设置超时，以便定期检查running状态
                        self.server_socket.settimeout(1.0)
                        client_socket, client_addr = self.server_socket.accept()
                        logger.info(f"接受到来自 {client_addr} 的连接")
                        
                        # 为每个客户端连接创建一个新线程处理转发
                        client_thread = threading.Thread(
                            target=self.handle_client, 
                            args=(client_socket, client_addr)
                        )
                        client_thread.daemon = True
                        client_thread.start()
                        self.threads.append(client_thread)
                    except socket.timeout:
                        continue  # 超时，继续检查running状态
                    except OSError as e:
                        # 处理套接字相关错误，特别是WinError 10038（在非套接字上操作）
                        if self.running:  # 只有在仍在运行时才记录错误
                            # 检查是否是非套接字操作错误
                            if hasattr(e, 'winerror') and e.winerror == 10038:
                                logger.debug(f"套接字已关闭，停止接受连接")
                            else:
                                logger.error(f"接受连接时发生套接字错误: {str(e)}")
                        break
                    except Exception as e:
                        if self.running:
                            logger.error(f"接受连接时发生未知错误: {str(e)}")
                        break
                    
            except KeyboardInterrupt:
                logger.info("接收到中断信号，正在停止服务...")
            finally:
                self.stop()
                
        except Exception as e:
            logger.error(f"启动服务失败: {str(e)}")
            self.stop()
            raise
    
    def start_async(self):
        """异步启动转发服务"""
        if not self.server_thread or not self.server_thread.is_alive():
            self.server_thread = threading.Thread(target=self.start)
            self.server_thread.daemon = True
            self.server_thread.start()
            return True
        return False
    
    def handle_client(self, client_socket, client_addr):
        """处理客户端连接并进行数据转发
        
        Args:
            client_socket: 客户端套接字
            client_addr: 客户端地址
        """
        remote_socket = None
        try:
            # 连接到远程服务器（本地127.0.0.1的服务）
            remote_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            remote_socket.connect((self.remote_ip, self.remote_port))
            logger.info(f"已连接到目标服务器 {self.remote_ip}:{self.remote_port}")
            
            # 创建两个线程分别处理双向数据流
            client_to_remote = threading.Thread(
                target=self.forward_data, 
                args=(client_socket, remote_socket, "客户端->目标")
            )
            remote_to_client = threading.Thread(
                target=self.forward_data, 
                args=(remote_socket, client_socket, "目标->客户端")
            )
            
            client_to_remote.daemon = True
            remote_to_client.daemon = True
            
            client_to_remote.start()
            remote_to_client.start()
            
            # 等待任一方向的数据传输结束
            client_to_remote.join()
            remote_to_client.join()
            
        except Exception as e:
            logger.error(f"处理连接时出错 {client_addr}: {str(e)}")
        finally:
            # 确保两个套接字都被关闭
            self.close_socket(client_socket)
            self.close_socket(remote_socket)
            logger.info(f"连接已关闭 {client_addr}")
    
    def forward_data(self, source, destination, direction):
        """在两个套接字之间转发数据
        
        Args:
            source: 源套接字
            destination: 目标套接字
            direction: 转发方向描述
        """
        try:
            while self.running:
                # 设置超时，以便定期检查running状态
                try:
                    source.settimeout(1.0)
                    # 从源读取数据
                    data = source.recv(self.buffer_size)
                    if not data:
                        break  # 连接关闭
                    
                    # 发送数据到目标
                    destination.sendall(data)
                    # 可以取消下面这行的注释以启用详细日志
                    # logger.debug(f"{direction} 转发 {len(data)} 字节")
                except socket.timeout:
                    continue  # 超时，继续检查running状态
                except OSError as e:
                    # 处理套接字相关错误
                    if self.running:
                        # 忽略正常关闭和非套接字错误
                        if hasattr(e, 'winerror') and e.winerror == 10038:
                            logger.debug(f"{direction} 套接字已关闭，停止转发")
                        elif isinstance(e, BrokenPipeError) or "远程主机强迫关闭了一个现有的连接" in str(e):
                            pass  # 这些是正常关闭错误，不记录
                        else:
                            logger.error(f"{direction} 数据转发套接字错误: {str(e)}")
                    break
                except Exception as e:
                    if self.running:
                        logger.error(f"{direction} 数据转发未知错误: {str(e)}")
                    break
        except Exception as e:
            # 最外层异常捕获，确保不会崩溃
            logger.error(f"{direction} 转发线程异常: {str(e)}")
    
    def stop(self):
        """停止转发服务"""
        self.running = False
        
        # 安全地关闭服务器套接字
        if hasattr(self, 'server_socket') and self.server_socket:
            try:
                self.close_socket(self.server_socket)
                self.server_socket = None  # 标记为None，避免后续使用
            except Exception as e:
                logger.debug(f"关闭服务器套接字时出错: {str(e)}")
        
        # 等待所有线程结束
        if hasattr(self, 'threads'):
            for thread in self.threads:
                if thread.is_alive():
                    try:
                        thread.join(0.5)  # 设置超时，避免卡死
                    except Exception as e:
                        logger.debug(f"等待线程结束时出错: {str(e)}")
            # 清空线程列表
            self.threads.clear()
        
        logger.info("端口转发服务已停止")
    
    @staticmethod
    def close_socket(sock):
        """安全关闭套接字"""
        if sock:
            try:
                # 尝试关闭套接字的写入端，允许读取端完成
                try:
                    sock.shutdown(socket.SHUT_WR)
                except OSError:
                    # 如果已经关闭或出错，继续执行
                    pass
                # 等待一小段时间让数据完成传输
                time.sleep(0.01)
                # 完全关闭套接字
                sock.close()
            except Exception as e:
                # 捕获所有异常，确保不会抛出错误
                logger.debug(f"关闭套接字时发生异常: {str(e)}")
                # 最后尝试简单关闭
                try:
                    sock.close()
                except:
                    pass


class PortForwarderApp:
    def __init__(self, root):
        self.root = root
        self.root.title("端口映射工具")
        self.root.geometry("700x500")
        # 禁用窗口最大化，设置为(False, False)表示完全禁止调整大小
        self.root.resizable(False, False)
        
        # 设置窗口图标为Tk内置的设置图标
        self.set_window_icon()
        
        # 确保窗口在设置位置前已经计算了大小
        self.root.update_idletasks()
        
        # 获取屏幕宽度和高度
        screen_width = self.root.winfo_screenwidth()
        screen_height = self.root.winfo_screenheight()
        
        # 获取窗口宽度和高度
        window_width = self.root.winfo_width()
        window_height = self.root.winfo_height()
        
        # 计算窗口居中的坐标
        x = (screen_width - window_width) // 2
        y = (screen_height - window_height) // 2
        
        # 设置窗口位置
        self.root.geometry(f"+{x}+{y}")
        
        # 设置中文字体支持
        self.set_fonts()
        
        # 创建日志队列和处理器
        self.log_queue = queue.Queue()
        self.queue_handler = QueueHandler(self.log_queue)
        self.queue_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
        logger.addHandler(self.queue_handler)
        
        # 创建PortForwarder实例
        self.forwarder = None
        
        # 创建界面
        self.create_widgets()
        
        # 加载上一次的配置
        self.load_config()
        
        # 启动日志更新线程
        self.update_log_thread = threading.Thread(target=self.update_logs, daemon=True)
        self.update_log_thread.start()
        
        # 设置窗口关闭事件
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        
        # 初始化系统托盘
        self.tray_icon = None
        self.create_system_tray()
    
    def set_window_icon(self):
        """设置窗口图标为Tk内置的推荐图标"""
        try:
            # 方法1: 使用Tk的默认图标机制
            if sys.platform.startswith('win'):
                self.root.iconbitmap(default='')
            else:
                self.root.iconbitmap('')
        except Exception as general_err:
            # 最后备选方案：完全依赖Tkinter的默认行为
            try:
                print(f"设置Tk内置图标的过程中发生错误: {str(general_err)}")
                # 清除所有可能的图标设置，让系统使用默认图标
                if hasattr(self.root, 'iconphoto'):
                    self.root.iconphoto(False) 
            except Exception as fallback_err:
                print(f"所有图标设置方法都失败: {str(fallback_err)}")
    
    def set_fonts(self):
        """设置字体配置"""
        if sys.platform == 'win32':
            # Windows平台字体
            self.font_family = 'Microsoft YaHei'
        else:
            # 其他平台字体
            self.font_family = 'SimHei'
        
        self.normal_font = (self.font_family, 10)
        self.bold_font = (self.font_family, 10, 'bold')
    
    def create_widgets(self):
        """创建GUI组件"""
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建配置框架
        config_frame = ttk.LabelFrame(main_frame, text="端口映射配置", padding="10")
        config_frame.pack(fill=tk.X, pady=(0, 10))
        
        # 监听IP
        ttk.Label(config_frame, text="监听IP:", font=self.normal_font).grid(row=0, column=0, sticky=tk.W, pady=5)
        self.listen_ip_var = tk.StringVar(value="0.0.0.0")
        ttk.Entry(config_frame, textvariable=self.listen_ip_var, width=20, font=self.normal_font).grid(row=0, column=1, padx=5, pady=5)
        
        # 监听端口
        ttk.Label(config_frame, text="监听端口:", font=self.normal_font).grid(row=0, column=2, sticky=tk.W, pady=5)
        self.listen_port_var = tk.StringVar(value="9000")
        ttk.Entry(config_frame, textvariable=self.listen_port_var, width=10, font=self.normal_font).grid(row=0, column=3, padx=5, pady=5)
        
        # 目标IP
        ttk.Label(config_frame, text="目标IP:", font=self.normal_font).grid(row=1, column=0, sticky=tk.W, pady=5)
        self.target_ip_var = tk.StringVar(value="127.0.0.1")
        ttk.Entry(config_frame, textvariable=self.target_ip_var, width=20, font=self.normal_font).grid(row=1, column=1, padx=5, pady=5)
        
        # 目标端口
        ttk.Label(config_frame, text="目标端口:", font=self.normal_font).grid(row=1, column=2, sticky=tk.W, pady=5)
        self.target_port_var = tk.StringVar(value="8080")
        ttk.Entry(config_frame, textvariable=self.target_port_var, width=10, font=self.normal_font).grid(row=1, column=3, padx=5, pady=5)
        
        # 缓冲区大小
        ttk.Label(config_frame, text="缓冲区大小:", font=self.normal_font).grid(row=2, column=0, sticky=tk.W, pady=5)
        self.buffer_size_var = tk.StringVar(value="4096")
        ttk.Entry(config_frame, textvariable=self.buffer_size_var, width=10, font=self.normal_font).grid(row=2, column=1, padx=5, pady=5)
        
        # 日志级别
        ttk.Label(config_frame, text="日志级别:", font=self.normal_font).grid(row=2, column=2, sticky=tk.W, pady=5)
        self.log_level_var = tk.StringVar(value="INFO")
        log_level_combo = ttk.Combobox(config_frame, textvariable=self.log_level_var, values=["INFO", "DEBUG"], width=8, font=self.normal_font)
        log_level_combo.grid(row=2, column=3, padx=5, pady=5)
        log_level_combo.bind("<<ComboboxSelected>>", self.on_log_level_change)
        
        # 启动/停止按钮
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=(0, 10))
        
        self.start_button = ttk.Button(button_frame, text="启动转发", command=self.start_forwarding, width=15)
        self.start_button.pack(side=tk.LEFT, padx=5)
        
        self.stop_button = ttk.Button(button_frame, text="停止转发", command=self.stop_forwarding, width=15, state=tk.DISABLED)
        self.stop_button.pack(side=tk.LEFT, padx=5)
        
        self.status_var = tk.StringVar(value="就绪")
        ttk.Label(button_frame, textvariable=self.status_var, font=self.bold_font).pack(side=tk.RIGHT, padx=10)
        
        # 日志显示区域
        log_frame = ttk.LabelFrame(main_frame, text="日志输出", padding="10")
        log_frame.pack(fill=tk.BOTH, expand=True)
        
        self.log_text = scrolledtext.ScrolledText(log_frame, wrap=tk.WORD, font=("Consolas", 9))
        self.log_text.pack(fill=tk.BOTH, expand=True)
        self.log_text.config(state=tk.DISABLED)
    
    def on_log_level_change(self, event=None):
        """处理日志级别变更"""
        level = self.log_level_var.get()
        if level == "DEBUG":
            logger.setLevel(logging.DEBUG)
        else:
            logger.setLevel(logging.INFO)
    
    def start_forwarding(self):
        """启动端口转发"""
        try:
            # 获取配置参数
            listen_ip = self.listen_ip_var.get().strip()
            target_ip = self.target_ip_var.get().strip()
            
            # 验证端口
            try:
                listen_port = int(self.listen_port_var.get().strip())
                target_port = int(self.target_port_var.get().strip())
                buffer_size = int(self.buffer_size_var.get().strip())
                
                if not (0 <= listen_port <= 65535 and 0 <= target_port <= 65535):
                    raise ValueError("端口号必须在0-65535范围内")
                
                if buffer_size <= 0:
                    raise ValueError("缓冲区大小必须大于0")
            except ValueError as e:
                self.show_error(f"参数错误: {str(e)}")
                return
            
            # 保存当前配置
            self.save_config()
            
            # 停止现有转发（如果有）
            self.stop_forwarding()
            
            # 创建并启动转发器
            self.forwarder = PortForwarder(
                local_ip=listen_ip,
                local_port=listen_port,
                remote_ip=target_ip,
                remote_port=target_port,
                buffer_size=buffer_size
            )
            
            # 更新UI状态
            self.start_button.config(state=tk.DISABLED)
            self.stop_button.config(state=tk.NORMAL)
            self.status_var.set("正在启动...")
            
            # 在单独的线程中启动转发服务
            def start_forwarder_safe():
                try:
                    self.forwarder.start_async()
                    # 延迟一点时间检查是否启动成功
                    time.sleep(0.5)
                    self.root.after(0, lambda: self.status_var.set("运行中"))
                    # 刷新托盘菜单状态
                    self.root.after(0, self.update_tray_menu)
                except Exception as e:
                    self.root.after(0, lambda: self.show_error(f"启动失败: {str(e)}"))
                    self.root.after(0, self.reset_ui_state)
            
            threading.Thread(target=start_forwarder_safe, daemon=True).start()
            
        except Exception as e:
            self.show_error(f"启动时出错: {str(e)}")
            self.reset_ui_state()
    
    def stop_forwarding(self):
        """停止端口转发"""
        if self.forwarder:
            try:
                self.forwarder.stop()
                self.forwarder = None
            except Exception as e:
                logger.error(f"停止转发时出错: {str(e)}")
        self.reset_ui_state()
        # 刷新托盘菜单状态
        self.update_tray_menu()
    
    def reset_ui_state(self):
        """重置UI状态"""
        self.start_button.config(state=tk.NORMAL)
        self.stop_button.config(state=tk.DISABLED)
        self.status_var.set("就绪")
    
    def update_tray_menu(self):
        """更新托盘菜单状态"""
        if HAS_PYSTRAY and self.tray_icon:
            try:
                # 刷新托盘图标，这会导致菜单状态重新计算
                self.tray_icon.update_menu()
            except Exception as e:
                logger.debug(f"更新托盘菜单时出错: {str(e)}")
    
    def show_error(self, message):
        """显示错误消息"""
        logger.error(message)
        # 可以添加一个消息框显示错误
        # tk.messagebox.showerror("错误", message)
    
    def update_logs(self):
        """更新日志显示"""
        while True:
            try:
                # 尝试从队列获取日志消息
                record = self.log_queue.get(block=False)
                self.root.after(0, lambda msg=record: self.append_log(msg))
            except queue.Empty:
                # 队列为空，等待一会儿
                time.sleep(0.1)
    
    def append_log(self, message):
        """向日志文本框添加消息"""
        self.log_text.config(state=tk.NORMAL)
        self.log_text.insert(tk.END, message + "\n")
        # 自动滚动到最后一行
        self.log_text.see(tk.END)
        self.log_text.config(state=tk.DISABLED)
    
    def on_closing(self):
        """处理窗口关闭事件 - 最小化到托盘而不是退出"""
        # 隐藏窗口而不是直接销毁
        self.root.withdraw()
        if self.tray_icon:
            logger.info("程序已最小化到系统托盘")
    
    def create_system_tray(self):
        """创建系统托盘图标和菜单"""
        if not HAS_PYSTRAY:
            logger.warning("系统托盘功能不可用，需要安装pystray和pillow库")
            return
        
        try:
            # 创建一个简单的托盘图标
            def create_tray_image(width=64, height=64, color1='#000080', color2='#FFFFFF'):
                # 创建图像
                image = Image.new('RGB', (width, height), color1)
                draw = ImageDraw.Draw(image)
                
                # 绘制一个简单的设置图标
                center_x, center_y = width // 2, height // 2
                radius = min(width, height) // 3
                
                # 绘制中心圆
                draw.ellipse((center_x - radius, center_y - radius, 
                             center_x + radius, center_y + radius), fill=color2)
                
                # 绘制设置符号
                line_length = radius * 2 // 3
                # 水平线
                draw.line((center_x - line_length, center_y, 
                           center_x + line_length, center_y), fill=color1, width=3)
                # 垂直线
                draw.line((center_x, center_y - line_length, 
                           center_x, center_y + line_length), fill=color1, width=3)
                
                return image
            
            # 创建托盘菜单
            menu = TrayMenu(
                TrayMenuItem("显示窗口", self.show_window, default=True),
                TrayMenuItem("连接", self.tray_connect, enabled=lambda item: not self.forwarder or (hasattr(self.forwarder, 'running') and not self.forwarder.running)),
                TrayMenuItem("断开", self.tray_disconnect, enabled=lambda item: self.forwarder and hasattr(self.forwarder, 'running') and self.forwarder.running),
                TrayMenuItem("退出", self.tray_exit)
            )
            
            # 创建托盘图标
            self.tray_icon = TrayIcon("端口映射工具", create_tray_image(), "端口映射工具", menu)
            
            # 在单独的线程中运行托盘
            self.tray_thread = threading.Thread(target=self.tray_icon.run, daemon=True)
            self.tray_thread.start()
            
            logger.info("系统托盘已创建")
            
        except Exception as e:
            logger.error(f"创建系统托盘时出错: {str(e)}")
            self.tray_icon = None
    
    def show_window(self):
        """从托盘中显示窗口"""
        self.root.deiconify()
        self.root.lift()  # 窗口置顶
        logger.info("窗口已显示")
    
    def tray_connect(self):
        """从托盘菜单启动连接"""
        # 保存当前配置
        self.save_config()
        # 在主线程中执行GUI操作
        self.root.after(0, self.start_forwarding)
    
    def tray_disconnect(self):
        """从托盘菜单断开连接"""
        # 在主线程中执行GUI操作
        self.root.after(0, self.stop_forwarding)
    
    def get_config_file_path(self):
        """获取配置文件路径"""
        # 在Windows上使用AppData目录，在其他系统上使用用户目录
        if sys.platform == 'win32':
            config_dir = os.path.join(os.environ['APPDATA'], 'HonghaiProxy')
        else:
            config_dir = os.path.expanduser('~/.honghaiproxy')
        
        # 确保配置目录存在
        os.makedirs(config_dir, exist_ok=True)
        
        # 返回配置文件路径
        return os.path.join(config_dir, 'config.json')
    
    def save_config(self):
        """保存当前配置到文件"""
        try:
            config = {
                'listen_ip': self.listen_ip_var.get(),
                'listen_port': self.listen_port_var.get(),
                'target_ip': self.target_ip_var.get(),
                'target_port': self.target_port_var.get(),
                'buffer_size': self.buffer_size_var.get(),
                'log_level': self.log_level_var.get()
            }
            
            config_path = self.get_config_file_path()
            with open(config_path, 'w', encoding='utf-8') as f:
                json.dump(config, f, ensure_ascii=False, indent=2)
            
            logger.debug(f"配置已保存到: {config_path}")
        except Exception as e:
            logger.error(f"保存配置失败: {str(e)}")
    
    def load_config(self):
        """从文件加载配置"""
        try:
            config_path = self.get_config_file_path()
            if os.path.exists(config_path):
                with open(config_path, 'r', encoding='utf-8') as f:
                    config = json.load(f)
                
                # 加载配置到UI变量
                if 'listen_ip' in config:
                    self.listen_ip_var.set(config['listen_ip'])
                if 'listen_port' in config:
                    self.listen_port_var.set(config['listen_port'])
                if 'target_ip' in config:
                    self.target_ip_var.set(config['target_ip'])
                if 'target_port' in config:
                    self.target_port_var.set(config['target_port'])
                if 'buffer_size' in config:
                    self.buffer_size_var.set(config['buffer_size'])
                if 'log_level' in config:
                    self.log_level_var.set(config['log_level'])
                    # 同时更新日志级别
                    self.on_log_level_change()
                
                logger.info(f"已加载配置文件: {config_path}")
            else:
                logger.debug(f"配置文件不存在，使用默认值: {config_path}")
        except Exception as e:
            logger.error(f"加载配置失败: {str(e)}")
    
    def on_closing(self):
        """处理窗口关闭事件 - 最小化到托盘而不是退出"""
        # 隐藏窗口而不是直接销毁
        self.root.withdraw()
        if self.tray_icon:
            logger.info("程序已最小化到系统托盘")
    
    def tray_exit(self):
        """从托盘菜单退出程序"""
        try:
            # 保存当前配置
            self.save_config()
            
            # 停止转发
            self.stop_forwarding()
            
            # 停止托盘
            if self.tray_icon:
                self.tray_icon.stop()
                
            # 关闭窗口
            self.root.destroy()
            
            logger.info("程序已退出")
            
        except Exception as e:
            logger.error(f"退出程序时出错: {str(e)}")
            # 确保程序能退出
            sys.exit(1)


def main_gui():
    """启动GUI应用"""
    root = tk.Tk()
    app = PortForwarderApp(root)
    root.mainloop()


def main_cli():
    """命令行模式主函数"""
    parser = argparse.ArgumentParser(description='端口映射工具 - 将本地服务端口映射到可被其他IP访问的端口')
    parser.add_argument('-l', '--listen-port', type=int, required=True, help='监听端口（对外提供服务的端口）')
    parser.add_argument('-t', '--target-port', type=int, required=True, help='目标端口（本地127.0.0.1上的服务端口）')
    parser.add_argument('-i', '--listen-ip', type=str, default='0.0.0.0', help='监听IP地址，默认0.0.0.0（所有接口）')
    parser.add_argument('-r', '--target-ip', type=str, default='127.0.0.1', help='目标IP地址，默认127.0.0.1')
    parser.add_argument('-b', '--buffer-size', type=int, default=4096, help='缓冲区大小，默认4096')
    parser.add_argument('-v', '--verbose', action='store_true', help='启用详细日志')
    parser.add_argument('--cli', action='store_true', help='使用命令行模式')
    
    args = parser.parse_args()
    
    # 如果启用了详细日志，设置日志级别为DEBUG
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    
    # 验证端口范围
    if not (0 <= args.listen_port <= 65535 and 0 <= args.target_port <= 65535):
        logger.error("端口号必须在0-65535范围内")
        return 1
    
    # 启动端口转发器
    forwarder = PortForwarder(
        local_ip=args.listen_ip,
        local_port=args.listen_port,
        remote_ip=args.target_ip,
        remote_port=args.target_port,
        buffer_size=args.buffer_size
    )
    
    try:
        forwarder.start()
    except KeyboardInterrupt:
        logger.info("用户中断，程序退出")
    except Exception as e:
        logger.error(f"程序异常: {str(e)}")
        return 1
    
    return 0


def main():
    """主入口函数"""
    # 检查是否有命令行参数
    if len(sys.argv) > 1:
        # 如果指定了--cli参数或者有其他参数，使用命令行模式
        if '--cli' in sys.argv or any(arg.startswith('-') for arg in sys.argv[1:]):
            return main_cli()
    
    # 默认使用GUI模式
    try:
        main_gui()
    except Exception as e:
        print(f"启动GUI失败: {str(e)}")
        print("尝试使用命令行模式...")
        return main_cli()


if __name__ == '__main__':
    sys.exit(main())